#!/home/wcaarls/compile/bin/python

import socket, os, thread, sys, random, traceback, time

def reserve():
  # Clear garbage
  lock.acquire()
  
  while 1:
    # Broadcast
    r = "%05d" % (random.randint(0, 65536))
    
    udp.sendto("?" + r, ('<broadcast>', 8533))
  
    # Return address of first return packet
    udp.settimeout(0.25)
    try:
      while 1:
        (data, address) = udp.recvfrom(6)
        if data == "!" + r:
          break
      if data == "!" + r:
        break
    except:
      pass

  conn = socket.socket(socket.AF_INET, socket.SOCK_STREAM)  
  conn.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
  conn.settimeout(10)
  conn.bind(('', 0))

  conn.connect(address)
  conn.settimeout(120);

  while 1:
    try:
      c = conn.recv(1)
      if c != '*':
        raise ValueError(c)
      break
    except e:
      print >>stderr, "Caught " + str(e)
      time.sleep(0.1)
      pass

  lock.release()
  
  return conn

def start(name, wait=False, conn=None):
  if conn == None:
    conn = __reserve()

  # Send program and parameters
  conn.send(name)
  
  try:
    c = conn.recv(1)
    if c != '+':
      raise ValueError(c)
  except:
    print "jobfarm.start: " + name + " failed to start"
    print "jobfarm.start: at " + str(conn.getpeername())
    traceback.print_exc()
    raise RuntimeError

  if wait:  
    # Wait for completion
    #print "% Awaiting completion of " + name
    try:
      c = conn.recv(1)
      if c != '.':
        raise ValueError(c)
    except:
      print "jobfarm.start: " + name + " failed to complete"
      print "jobfarm.start: at " + str(conn.getpeername())
      traceback.print_exc()
      raise RuntimeError
      
  conn.close()

def suspend():
    udp = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    udp.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    udp.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
    udp.bind(('', 0))
    udp.sendto("SUSPND", ('<broadcast>', 8533))

def resume():
    udp = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    udp.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    udp.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
    udp.bind(('', 0))
    udp.sendto("RESUME", ('<broadcast>', 8533))

def __spawn(conn):
    global threads

    conn.send('*')

    # Get program name
    print "Awaiting process parameters"
    
    conn.settimeout(10)
    data = conn.recv(1024)
    
    if data and data != '':
      print "Executing '" + data + "'"
     
      conn.send('+')
    
      # Start program
      T = time.time()
      if os.system(data):
        cleanexit = False
      else:
        cleanexit = True
      T = time.time() - T
    
      if cleanexit:
        print "Child terminated after " + str(T) + " seconds"
      else:
        print "Child terminated abnormally after " + str(T) + " seconds"
        print "Command was '" + data + "'"
        raise RuntimeError
    
      # Signal completion
      conn.send('.')
    else:
      print "Job aborted"
    
    conn.close()
    
    lock.acquire()
    
    threads = threads - 1
    if threads:
      print str(threads) + " jobs remaining"
      
    lock.release()

def __accept():
    global threads
    
    nr = 0

    tcp = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    tcp.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    tcp.bind(('', 8533))
    tcp.listen(5)

    while 1:
      (conn, address) = tcp.accept()
      
      print "Incoming connection (" + str(nr) + ") from " + str(address)
      
      nr = nr + 1
      
      lock.acquire()
      threads = threads + 1
      lock.release()
      
      thread.start_new_thread(__spawn, (conn,))
    
def slave():
  # Slave
  print "SmartCam jobfarm slave"
  
  try:
    arg = sys.argv[1]
  except:
    arg = "1"

  try:
    maxthreads = int(arg)
    print "Servicing " + str(maxthreads) + " jobs at once"
  except:
    if arg == "suspend":
      print "Broadcasting jobfarm suspension"
      suspend()
    elif arg == "resume":
      print "Broadcasting jobfarm resumption"
      resume()
    else:
      print "Unknown argument " + arg
    sys.exit(1)
  
  # Bind sockets
  udp = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
  udp.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
  udp.bind(('', 8533))
  
  thread.start_new_thread(__accept, ())

  print "Listening on port 8533"
  
  hold = 0
  while 1:
    # Listen for UDP broadcasts
    (data, address) = udp.recvfrom(6)

    if data[0] == "?":
      lock.acquire()
      if threads < maxthreads and not hold:
        print "Probe " + data + " from " + str(address)

        # Reply
        udp.sendto("!" + data[1:], address)
      lock.release()
    elif data == "SUSPND":
      print "Suspended"
      hold = 1
    elif data == "RESUME":
      print "Resumed"
      hold = 0
    else:
      print "Unknown packet '" + data + "'"

threads = 0

# Bind UDP socket
udp = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
udp.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
udp.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
udp.bind(('', 0))

lock = thread.allocate_lock()

if __name__ == "__main__":
  slave()
  